import idc
import idaapi
import idautils

class RISC(object):

	CODE = 2
	DATA = 3
	SEARCH_DEPTH = 25

	ARCHITECTURES = {
		'mips'	: {
				'arguments'	: ['$a0', '$a1', '$a2', '$a3'],
				'savedregs'	: ['$s0', '$s1', '$s2', '$s3', '$s4', '$s5', '$s6', '$s7', '$fp', '$gp', '$ra'],
		},
		'arm'	: {
				'arguments'	: ['R0', 'R1', 'R2', 'R3'],
				'savedregs'	: [],
		}
	}

	def __init__(self):
		# This assumes that you will only ever analyze 32-bit binaries in 32-bit IDA...
		if idc.__EA64__:
			self.INSIZE = 8
		else:
			self.INSIZE = 4

		self.arch = self.ARCHITECTURES['mips']
		self.processor = idc.GetShortPrm(idc.INF_PROCNAME).lower()

		for (name, info) in self.ARCHITECTURES.iteritems():
			if self.processor.startswith(name.lower()):
				self.arch = self.ARCHITECTURES[name]
				break

	# Get the start of the specified segment type (2 == code, 3 == data)
	def get_start_ea(self, attr):
		ea = idc.BADADDR
		seg = idc.FirstSeg()

		while seg != idc.BADADDR:
			if idc.GetSegmentAttr(seg, idc.SEGATTR_TYPE) == attr:
				ea = seg
				break
			else:
				seg = idc.NextSeg(seg)
	
		return ea

	# Creates ASCII strings and converts remaining data into DWORDS.
	def datify(self):
		n = 0
		ea = self.get_start_ea(self.DATA)

		print "\nLooking for possible strings starting at: %s:0x%X" % (idc.SegName(ea), ea)

		for s in idautils.Strings():
			if s.ea > ea:
				if not idc.isASCII(idc.GetFlags(s.ea)) and idc.MakeStr(s.ea, idc.BADADDR):
					n += 1

		print "Created %d new ASCII strings" % n
	
		print "Converting remaining data to DWORDs...",
	
		while ea != idc.BADADDR:
			flags = idc.GetFlags(ea)
		            
			if idc.isUnknown(flags) or idc.isByte(flags):
				idc.MakeDword(ea)
				idc.OpOff(ea, 0, 0)

			ea = idc.NextAddr(ea)

		print "done.\n"

	def codeify(self, ea=idc.BADADDR, force=False):
		func_count = 0
		code_count = 0

		if ea == idc.BADADDR:
			ea = self.get_start_ea(self.CODE)

		if not force and self.get_start_ea(self.DATA) == idc.BADADDR:
			print "WARNING: No data segments defined! I don't know where the code segment ends and the data segment begins."
	
		print "\nLooking for undefined code starting at: 0x%X...\n" % ea

		while ea != idc.BADADDR:
			try:
				if idc.GetSegmentAttr(ea, idc.SEGATTR_TYPE) == self.CODE:
					if idc.GetFunctionName(ea) != '':
						ea = idc.FindFuncEnd(ea)
						continue
					else:
						if idc.MakeFunction(ea):
							func_count += 1
						elif idc.MakeCode(ea):
							code_count += 1
			except:
				pass
            
			ea = idc.NextAddr(ea)
    
		print "\nCreated %d new functions and %d new code blocks\n" % (func_count, code_count)


	def find_main(self):
		done = False
		libc_main = ['__libc_main', '__uClibc_main']

		if idc.LocByName('main') == idc.BADADDR:
			for libc_main_name in libc_main:
				location = idc.LocByName(libc_main_name)
				if location != idc.BADADDR:
					break

			if location != idc.BADADDR:
				print "Looking for references to main()..."
				for xref in idautils.XrefsTo(location):
					if xref.type != 1:
						for i in range(1, self.SEARCH_DEPTH):
							ea = xref.frm - (i * self.INSIZE)
							if idc.GetOpnd(ea, 0) == self.arch['arguments'][0]:
								main_address = idc.LocByName(idc.GetOpnd(ea, 1))
								if main_address != idc.BADADDR:
									idc.MakeName(main_address, 'main')
									idc.Jump(main_address)
									done = True
									break
					if done:
						break


	def name_saved_registers(self):
		func = None

		if self.processor.startswith('mips'):
			func = self.mipsvars

		if func:
			print "Naming saved register locations...",
			func()
			print "done."

	def mipsvars(self):
		for ea in idautils.Functions():
			mea = ea
			named_regs = []
			last_iteration = False

			while mea < (ea + (self.INSIZE * self.SEARCH_DEPTH)):
				mnem = idc.GetMnem(mea)

				if mnem == 'sw':
					reg = idc.GetOpnd(mea, 0)
					dst = idc.GetOpnd(mea, 1)
	
					if reg in self.arch['savedregs'] and reg not in named_regs and dst.endswith('($sp)') and 'var_' in dst:
						offset = int(dst.split('var_')[1].split('(')[0], 16)
						idc.MakeLocal(ea, idc.FindFuncEnd(ea), "[sp-%d]" % offset, "saved_%s" % reg[1:])
						named_regs.append(reg)
				
				if last_iteration:
					break
				elif mnem.startswith('j') or mnem.startswith('b'):
					last_iteration = True

				mea += self.INSIZE

class rischelper_t(idaapi.plugin_t):
	flags = 0
	comment = "RISC IDA Script Collection"
	help = ""
	wanted_name = "RISC IDA Script Collection"
	wanted_hotkey = ""

	def init(self):
		self.menu_context1 = idaapi.add_menu_item("Options/", "Fixup data/code", "Alt-3", 0, self.fix_code_data, (None,))
		if idc.GetShortPrm(idc.INF_PROCNAME).lower().startswith('mips'):
			self.menu_context2 = idaapi.add_menu_item("Options/", "Name saved registers", "Alt-4", 0, self.name_saved_registers, (None,))
		else:
			self.menu_context2 = None

		return idaapi.PLUGIN_KEEP

	def term(self):
		idaapi.del_menu_item(self.menu_context1)
		if self.menu_context2 is not None:
			idaapi.del_menu_item(self.menu_context2)
		return None

	def fix_code_data(self, arg):
		risc = RISC()
		risc.datify()
		risc.codeify()

	def name_saved_registers(self, arg):
		RISC().name_saved_registers()

	def run(self, arg):
		risc = RISC()
		risc.datify()
		risc.codeify()
		risc.name_saved_registers()
		risc.find_main()

def PLUGIN_ENTRY():
	return rischelper_t()

